import os
import json
from PIL import Image

def load_label_map(label_map_file):
    with open(label_map_file, 'r', encoding='utf-8') as f:
        label_map = json.load(f)
    reversed_label_map = {v: int(k) for k, v in label_map.items()}
    return reversed_label_map

def convert_to_coco_format(json_folder, image_folder, label_map_file, output_file):
    label_map = load_label_map(label_map_file)
    
    coco_format = {
        "images": [],
        "annotations": [],
        "categories": []
    }
    
    category_id_map = {}
    annotation_id = 1
    
    # Create categories
    for category, id in label_map.items():
        coco_format["categories"].append({
            "id": id,
            "name": category,
            "supercategory": "none"
        })
        category_id_map[category] = id
    
    # Process each JSON file
    for filename in os.listdir(json_folder):
        if filename.endswith('.json'):
            json_file = os.path.join(json_folder, filename)
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            image_filename = data['imagePath']
            image_path = os.path.join(image_folder, image_filename)
            with Image.open(image_path) as img:
                width, height = img.size
            
            image_id = len(coco_format["images"]) + 1
            coco_format["images"].append({
                "id": image_id,
                "file_name": image_filename,
                "height": height,
                "width": width
            })
            
            for shape in data['shapes']:
                category = shape['label']
                points = shape['points']
                xmin = float(points[0][0])
                ymin = float(points[0][1])
                xmax = float(points[1][0])
                ymax = float(points[1][1])
                if category == 'person':
                    category = 'overhead view worker'
                label = label_map.get(category, -1)
                if label == -1:
                    print(f"Warning: Category '{category}' not found in label map.")
                    continue
                
                bbox = [xmin, ymin, xmax - xmin, ymax - ymin]
                area = (xmax - xmin) * (ymax - ymin)
                
                coco_format["annotations"].append({
                    "id": annotation_id,
                    "image_id": image_id,
                    "category_id": label,
                    "bbox": bbox,
                    "area": area,
                    "iscrowd": 0
                })
                annotation_id += 1
    
    # Save to output file
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(coco_format, f, ensure_ascii=False, indent=4)

# 示例用法
json_folder = '/media/wmq/9E70720F6F1BDCF7/DATA/meishan_newdata_wmq_zmh_wzp/annotations'  # 替换为你的JSON文件夹路径
image_folder = '/media/wmq/9E70720F6F1BDCF7/DATA/meishan_newdata_wmq_zmh_wzp/images'  # 替换为你的图片文件夹路径
output_file = 'Meishan/annotation/val_coco.json'
label_map_file = 'Meishan/annotation/label_map.json'  # 替换为你的label_map.json文件路径
convert_to_coco_format(json_folder, image_folder, label_map_file, output_file)